/***************************************************************************
 *   Copyright (C) 2010 by programmeboy                                    *
 *   programmeboy@gmail.com                                                *
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU General Public License     *
 *   along with this program; if not, write to the                         *
 *   Free Software Foundation, Inc.                                        *
 *   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             * 
 ***************************************************************************/

#include "stdafx.h"
#include <tchar.h>
#include <Windows.h>
#include <Shlwapi.h>

#pragma comment(lib, "Shlwapi.lib")
#pragma warning(disable:4996)

#include "Utils.h"

CUtils::CUtils(void)
{
}

CUtils::~CUtils(void)
{
}

#define UNCPATH_PREFIX  _T("\\??\\UNC\\")
#define NTPATH_PREFIX   _T("\\??\\")
#define SYSTEMROOT_PREFIX   _T("\\SystemRoot\\")
#define WINDOWS_PREFIX  _T("\\WINDOWS\\")

BOOL 
CUtils::NtPathNameToDosPathName(
    IN OUT LPTSTR lpNtPathName, 
    IN ULONG Length
)
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
    //
    //  Is UNC Path?
    //
    
    if (0 == _tcsncmp(lpNtPathName, 
              UNCPATH_PREFIX, _tcslen(UNCPATH_PREFIX))){
        MoveMemory(lpNtPathName, lpNtPathName + _tcslen(UNCPATH_PREFIX), 
                   (Length - _tcslen(UNCPATH_PREFIX)) * sizeof(TCHAR));
    }else if (0 == _tcsncmp(lpNtPathName, 
                   NTPATH_PREFIX, _tcslen(NTPATH_PREFIX))){
       MoveMemory(lpNtPathName, lpNtPathName + _tcslen(NTPATH_PREFIX), 
                   (Length - _tcslen(NTPATH_PREFIX)) * sizeof(TCHAR));

    }else if (0 == _tcsncicmp(lpNtPathName, 
                    SYSTEMROOT_PREFIX, _tcslen(SYSTEMROOT_PREFIX))){
        TCHAR lpWindowsPath[MAX_PATH];
        GetWindowsDirectory(lpWindowsPath, MAX_PATH);

        MoveMemory(lpNtPathName, lpNtPathName + _tcslen(SYSTEMROOT_PREFIX), 
            Length - _tcslen(SYSTEMROOT_PREFIX));
        PathCombine(lpWindowsPath, lpWindowsPath, lpNtPathName);
        _tcscpy_s(lpNtPathName, Length, lpWindowsPath);

    }else if (0 == _tcsncmp(lpNtPathName, 
                   WINDOWS_PREFIX, _tcslen(WINDOWS_PREFIX))){
        TCHAR lpWindowsPath[MAX_PATH];
        GetWindowsDirectory(lpWindowsPath, MAX_PATH);

        MoveMemory(lpNtPathName, lpNtPathName + _tcslen(WINDOWS_PREFIX), 
            Length - _tcslen(WINDOWS_PREFIX));
        PathCombine(lpWindowsPath, lpWindowsPath, lpNtPathName);
        _tcscpy_s(lpNtPathName, Length, lpWindowsPath);
    }
    return TRUE;
}

BOOL 
CUtils::QuerySystemModuleInformation(
    IN LPCSTR ModuleName, 
    OUT PSYSTEM_MODULE ModuleInfo
)
{
    ULONG BufferLen;
    NTSTATUS ntStatus;
    PVOID Buffer;
    BOOL bRet;

    if (!ModuleInfo){
        return FALSE;
    }
    
    bRet = FALSE;

    //
    //  First call the function to get the size of buffer
    //

    ntStatus = ZwQuerySystemInformation( SystemModuleInformation,
        NULL,
        0,
        &BufferLen);
    if (ntStatus != STATUS_INFO_LENGTH_MISMATCH){
        return FALSE;
    }


    //
    //  Allocate the buffer for receive the information
    //

    if (!BufferLen){
        return FALSE;
    }

    Buffer = new CHAR[BufferLen];
    if (!Buffer){
        return FALSE;
    }

    //
    //  Call again
    //

    ntStatus = ZwQuerySystemInformation( SystemModuleInformation,
                                         Buffer,
                                         BufferLen,
                                         &BufferLen );
    if (NT_SUCCESS(ntStatus)){

        if (!ModuleName){

            //
            //  For kernel
            //

            *ModuleInfo = ((PSYSTEM_MODULE_INFORMATION)Buffer)->Modules[0];
            bRet = TRUE;
        }else{

            //
            //  another module.
            //

            for (ULONG i = 0; i < ((PSYSTEM_MODULE_INFORMATION)Buffer)->ModulesCount; i++ ){

                LPSTR ModuleNameTmp;

                ModuleNameTmp = ((PSYSTEM_MODULE_INFORMATION)Buffer)->Modules[i].ImageName + \
                    ((PSYSTEM_MODULE_INFORMATION)Buffer)->Modules[i].ModuleNameOffset;
                if (0 == lstrcmpiA(ModuleName, ModuleNameTmp)){
                    *ModuleInfo = ((PSYSTEM_MODULE_INFORMATION)Buffer)->Modules[i];
                    bRet = TRUE;
                    break;
                }
            }
        }
    }

    //
    //  free the buffer
    //

    delete Buffer;
    return bRet;
}

BOOL
CUtils::QueryKernelMoudleInformation(
    OUT PSYSTEM_MODULE ModuleInfo
)
{
    return QuerySystemModuleInformation(NULL, ModuleInfo);
}

#define IS_ADDRESS_IN_MODULE(_base, _size, _address)    \
                (((ULONG)_address >= (ULONG)_base) &&   \
                ((ULONG)_address <= ((ULONG)_base + (ULONG)_size)))

BOOL
CUtils::ConvertSysModPathToNormal(
    IN LPCSTR lpszSysModPath,
    OUT LPTSTR lpNormalPath,
    IN ULONG Length
)
{
    TCHAR szPathTmp[MAX_PATH] = {0};

	lstrcpyn(szPathTmp, (LPCTSTR)CA2T(lpszSysModPath), MAX_PATH);

    //
    //  the path is nt's style. so we must convert it
    //

    if(!CUtils::NtPathNameToDosPathName(szPathTmp, MAX_PATH))
        return FALSE;
    lstrcpy(lpNormalPath, szPathTmp);
    return TRUE;
}

BOOL 
CUtils::QueryModuleNameByAddress(
    IN ULONG Address,
    OUT LPTSTR ModulePath,
    IN ULONG Length
)
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
    ULONG BufferLen;
    NTSTATUS ntStatus;
    PVOID Buffer;

    if (!Address || !ModulePath){
        return FALSE;
    }

    //
    //  First call the function to get the size of buffer
    //

    ntStatus = ZwQuerySystemInformation( SystemModuleInformation,
                                         NULL,
                                         0,
                                         &BufferLen);
    if (ntStatus != STATUS_INFO_LENGTH_MISMATCH){
        return FALSE;
    }


    //
    //  Allocate the buffer for receive the information
    //

    if (!BufferLen){
        return FALSE;
    }

    Buffer = new CHAR[BufferLen];
    if (!Buffer){
        return FALSE;
    }

    //
    //  Call again
    //

    ntStatus = ZwQuerySystemInformation( SystemModuleInformation,
                                         Buffer,
                                         BufferLen,
                                         &BufferLen );
    if (NT_SUCCESS(ntStatus)){

        //
        //  the first is the kernel's information
        //

        for (ULONG i = 0; i < ((PSYSTEM_MODULE_INFORMATION)Buffer)->ModulesCount; i++ ){

            PSYSTEM_MODULE ModuleInfo;

            ModuleInfo = &((PSYSTEM_MODULE_INFORMATION)Buffer)->Modules[i];
            if (IS_ADDRESS_IN_MODULE(ModuleInfo->Base, ModuleInfo->Size, Address)){

				TCHAR szModTmp[MAX_PATH] = {0};
				LPCTSTR lpImageName = CA2T(ModuleInfo->ImageName);
				if (i == 0){
					GetSystemDirectory(szModTmp, MAX_PATH);
					PathAppend(szModTmp, PathFindFileName(lpImageName));
				}else{
					lstrcpyn(szModTmp, lpImageName, MAX_PATH);
				}
				

                //
                //  the path is nt's style. so we must convert it
                //

                CUtils::NtPathNameToDosPathName(szModTmp, MAX_PATH);
				
				//
				// copy to out 
				//
				
				lstrcpyn(ModulePath, szModTmp, Length);

                break;

            }
        }
    }

    //
    //  free the buffer
    //

    delete Buffer;
    return TRUE;
}

BOOL
CUtils::QueryKernelPath(
    OUT LPTSTR lpszPath,
    IN ULONG Length
)
{
    SYSTEM_MODULE KernelInfo;
    TCHAR szNormalPath[MAX_PATH] = {0};

    //
    //  first find the kernel image.
    //

    if(!QueryKernelMoudleInformation(&KernelInfo)){
        LoggingErrorString(_T("query kernel module error"));
        return FALSE;
    }

    //
    //  convert the image path to normal.
    //

    if (!ConvertSysModPathToNormal(KernelInfo.ImageName, szNormalPath, MAX_PATH)){
        LoggingErrorString(_T("ConvertSysModPathToNormal error"));
        return FALSE;
    }

    //
    // copy the output.
    //

    if ((ULONG)lstrlen(szNormalPath) > Length){
        return FALSE;
    }

    lstrcpy(lpszPath, szNormalPath);

    //
    // Is file exist?
    //

    if (!PathFileExists(szNormalPath)){
        return FALSE;
    }

    return TRUE;
}

ULONG
CUtils::QueryKernelBase(
    VOID
)
{
    SYSTEM_MODULE KernelInfo;

    //
    //  first find the kernel image.
    //

    if(!QueryKernelMoudleInformation(&KernelInfo)){

        //
        //  TODO Logging the error.
        //

        return FALSE;
    }

    return (ULONG)KernelInfo.Base;
}